import os
import telnetlib
import time
from datetime import datetime
import pandas as pd
from openpyxl.utils.dataframe import dataframe_to_rows
from pandas import DataFrame

from openpyxl import load_workbook, Workbook
from openpyxl.chart import LineChart, Reference


def get_system_usage(host, port, user, password, command):
    try:
        with telnetlib.Telnet(host, port, timeout=10) as tn:
            tn.read_until(b"login: ")
            tn.write(user + b"\n")
            tn.read_until(b"Password: ")
            tn.write(password + b"\n")
            tn.write(command + b"\n")
            time.sleep(1)  # 等待命令执行和输出

            output = tn.read_very_eager().decode('ascii')

            cpu_usage = memory_usage = None
            for line in output.splitlines():
                if "CPU_USAGE=" in line:
                    cpu_usage = float(line.split("=")[1])
                elif "MEM_USAGE=" in line:
                    memory_usage = float(line.split("=")[1])

            return cpu_usage, memory_usage
    except Exception as e:
        print(f"Error: {e}")
        return None, None


from openpyxl.styles import Font


def add_charts_and_formatting_to_sheet(filename, sheet_name, cpu_threshold=70, mem_threshold=80):
    book = load_workbook(filename)
    sheet = book[sheet_name]

    # 创建图表，参照之前的add_charts_to_sheet函数内容
    book = load_workbook(filename)
    sheet = book[sheet_name]

    # 创建CPU使用率图表
    cpu_chart = LineChart()
    cpu_chart.title = "CPU Usage"
    cpu_chart.style = 13
    cpu_chart.y_axis.title = 'CPU Usage (%)'
    cpu_chart.x_axis.title = 'Time'
    cpu_chart.width = 30  # 图表宽度
    cpu_chart.height = 10  # 图表高度

    # 创建内存使用率图表
    mem_chart = LineChart()
    mem_chart.title = "Memory Usage"
    mem_chart.style = 13
    mem_chart.y_axis.title = 'Memory Usage (%)'
    mem_chart.x_axis.title = 'Time'

    mem_chart.width = 30  # 图表宽度
    mem_chart.height = 10  # 图表高度

    # 添加CPU使用率数据到图表
    cpu_data = Reference(sheet, min_col=2, min_row=1, max_col=2, max_row=sheet.max_row)
    cpu_chart.add_data(cpu_data, titles_from_data=True)

    # 添加时间作为图表的X轴
    times = Reference(sheet, min_col=1, min_row=2, max_row=sheet.max_row)
    cpu_chart.set_categories(times)

    # 同理添加内存使用率数据到图表
    mem_data = Reference(sheet, min_col=3, min_row=1, max_col=3, max_row=sheet.max_row)
    mem_chart.add_data(mem_data, titles_from_data=True)
    mem_chart.set_categories(times)

    # 将图表添加到工作表
    sheet.add_chart(cpu_chart, "E2")
    sheet.add_chart(mem_chart, "E20")

    book.save(filename)

    # 请根据前面的示例添加创建和添加图表的代码

    # 字体样式：红色
    red_font = Font(color="FF0000")

    # 检查并格式化超过阈值的CPU使用率单元格
    for row in range(2, sheet.max_row + 1):  # 从第2行开始到最大行数
        cpu_cell = sheet.cell(row=row, column=2)  # CPU使用率在第二列
        if cpu_cell.value is not None and cpu_cell.value > cpu_threshold:
            cpu_cell.font = red_font

    # 检查并格式化超过阈值的内存使用率单元格
    for row in range(2, sheet.max_row + 1):
        mem_cell = sheet.cell(row=row, column=3)  # 内存使用率在第三列
        if mem_cell.value is not None and mem_cell.value > mem_threshold:
            mem_cell.font = red_font

    book.save(filename)


from openpyxl import load_workbook


def save_data_to_excel(data, filename, sheet_name):
    # 将数据转换为pandas数据帧
    df = DataFrame(data)

    # 检查文件是否存在来决定是加载还是创建工作簿
    if os.path.exists(filename):
        workbook = load_workbook(filename)
        # 如果存在同名的工作表，则先删除
        if sheet_name in workbook.sheetnames:
            del workbook[sheet_name]
        worksheet = workbook.create_sheet(sheet_name)
    else:
        workbook = Workbook()
        worksheet = workbook.active
        worksheet.title = sheet_name

    # 使用openpyxl将pandas数据帧转换为工作表行
    for r_idx, row in enumerate(dataframe_to_rows(df, index=False, header=True), 1):
        for c_idx, value in enumerate(row, 1):
            worksheet.cell(row=r_idx, column=c_idx, value=value)

    # 保存工作簿
    workbook.save(filename)

    # 保存完数据后，再单独处理图表和条件格式，因为pandas ExcelWriter提交(write)后无法添加
    add_charts_and_formatting_to_sheet(filename, sheet_name)


def main(device_name, device_ip, run_time):
    # 设置
    HOST = device_ip
    PORT = "23"
    USER = b"root\n"
    PASSWORD = b"vp202401\n"
    COMMAND = b"./tmp/get_system_usage.sh\n"
    data = []

    start_time = datetime.now()
    while (datetime.now() - start_time).seconds < run_time:  # 监控时间
        cpu_usage, memory_usage = get_system_usage(HOST, PORT, USER, PASSWORD, COMMAND)
        if cpu_usage is not None and memory_usage is not None:
            data.append({
                'Time': datetime.now().strftime(' %H:%M:%S'),
                'CPU_Usage': cpu_usage,
                'Memory_Usage': memory_usage
            })
        time.sleep(1)  # 每10秒采集一次数据

    save_data_to_excel(data, 'system_usage.xlsx', datetime.now().strftime(f'{device_name}_%Y-%m-%d'))


if __name__ == "__main__":
    main('NQ10-1', '192.168.2.98', 30)
